package org.ws.task;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.netty.channel.Channel;
import io.netty.util.concurrent.DefaultThreadFactory;
import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.ws.task.builder.RequestMessageBuilder;
import org.ws.task.event.EventListener;
import org.ws.task.event.PushItemEvent;
import org.ws.task.event.support.BootstrapEventHolder;
import org.ws.task.event.support.PushItemEventHolder;
import org.ws.task.event.support.startup.BootstrapEvent;
import org.ws.task.event.support.startup.BootstrapEventType;
import org.ws.task.item.ExecuteContext;
import org.ws.task.item.TaskInstance;
import org.ws.task.netty.handler.ActionHandlerRegistry;
import org.ws.task.netty.handler.CommitStatusHandler;
import org.ws.task.netty.handler.DefaultActionHandlerRegistry;
import org.ws.task.netty.handler.PullItemProviderHandler;
import org.ws.task.netty.handler.PushItemLeaderHandler;
import org.ws.task.repository.LeaderContextRepository;
import org.ws.task.repository.NodeMetadataRepository;
import org.ws.task.repository.TaskRepository;
import org.ws.task.model.NodeMetadata;
import org.ws.task.model.protocol.PushItemRequest;
import org.ws.task.model.TaskItem;
import org.ws.task.netty.NettyConfig;
import org.ws.task.netty.NettyServer;
import org.ws.task.netty.manager.ChannelConnectManager;
import org.ws.task.model.protocol.base.RequestAction;
import org.ws.task.model.protocol.base.RequestMessage;
import org.ws.task.model.protocol.base.ResponseMessage;
import org.ws.task.model.TaskProcessStatus;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Leader控制器
 * 数据同步
 * 数据管理
 * 任务调度管理
 * 统一分配管理
 */
public class LeaderController implements Bootable {

    private static Logger log = LoggerFactory.getLogger(LeaderController.class.getName());

    private final NettyServer nettyServer;
    private final ActionHandlerRegistry actionHandlerRegistry;
    // 拉数据处理
    private final PullItemProviderHandler pullItemProviderHandler;
    // 提交状态处理
    private final CommitStatusHandler commitStatusHandler;
    // 推数据处理
    private final PushItemLeaderHandler pushItemLeaderHandler;
    private final PushItemEventHandler pushItemEventHandler;
    private final ExecutorService actionHandlerExecutor;

    public LeaderController(NettyServer nettyServer, NettyConfig nettyConfig) {
        this.nettyServer = nettyServer;

        this.actionHandlerExecutor = Executors.newFixedThreadPool(nettyConfig.getExecutorThreadNum(),
                new DefaultThreadFactory("LeaderActionHandlerThread"));

        this.pullItemProviderHandler = new PullItemProviderHandler();
        this.pushItemLeaderHandler = new PushItemLeaderHandler();
        this.actionHandlerRegistry = new DefaultActionHandlerRegistry();
        this.commitStatusHandler = new CommitStatusHandler();
        this.pushItemEventHandler = new PushItemEventHandler(this.nettyServer.getChannelConnectManager());
    }

    @Override
    public void start() {
        this.actionHandlerRegistry.register(this.pullItemProviderHandler);
        this.actionHandlerRegistry.register(this.commitStatusHandler);
        this.actionHandlerRegistry.register(this.pushItemLeaderHandler);

        this.nettyServer.actionHandlerRegistry(this.actionHandlerRegistry,this.actionHandlerExecutor);

        PushItemEventHolder.getInstance().registerListener(this.pushItemEventHandler);

        log.info("Leader Controller start.");
    }

    @Override
    public void shutdown() {
        this.actionHandlerRegistry.clear();
        this.nettyServer.shutdown();
    }

    /**
     * 推送事件处理
     */
    class PushItemEventHandler implements EventListener<PushItemEvent> {

        private final ChannelConnectManager channelConnectManager;
        private final Map<String,AtomicBoolean> taskNameRunningMap;

        PushItemEventHandler(ChannelConnectManager channelConnectManager) {
            this.channelConnectManager = channelConnectManager;
            this.taskNameRunningMap = Maps.newConcurrentMap();
        }

        @Override
        public void onEvent(PushItemEvent event) {
            // 执行推送消息
            String taskName = event.getSource();
            if(!this.taskNameRunningMap.containsKey(taskName)){
                this.taskNameRunningMap.put(taskName,new AtomicBoolean(false));
            }
            if(this.taskNameRunningMap.get(taskName).get()){
                return;
            }
            this.taskNameRunningMap.get(taskName).set(true);
            TaskInstance taskInstance = TaskRepository.getInstance().getTaskInstance(taskName);
            long fetchTimeout = taskInstance.getTaskConfig().getFetchTimeout();
            int chunk = taskInstance.getTaskConfig().getChunk();
            List<NodeMetadata> workerNodes = NodeMetadataRepository.getInstance().getActiveWorkerNodes(taskName);
            try {
                if (CollectionUtils.isNotEmpty(workerNodes)) {
                    for (NodeMetadata node : workerNodes) {
                        String id = node.getId();
                        String address = node.getAddress();
                        ExecuteContext context = LeaderContextRepository.getInstance().getContext(taskName, id);
                        int i = 0;
                        List<TaskItem> list = Lists.newArrayList();
                        while (!context.isEmpty()) {
                            // 全部推送
                            TaskItem item = context.poll();
                            if(item != null){
                                list.add(item);
                                i++;
                                if(i % chunk == 0){
                                    pushAsync(address,taskName, list , context, fetchTimeout);
                                    list = Lists.newArrayList();
                                }
                            }
                        }
                        if(CollectionUtils.isNotEmpty(list)){
                            pushAsync(address,taskName, list , context, fetchTimeout);
                        }
                    }
                }
            }
            finally {
                this.taskNameRunningMap.get(taskName).set(false);
            }
        }

        private void pushAsync(String addr, String taskName, Collection<TaskItem> items, ExecuteContext context, long timeout){

            log.debug("pushAsync [{}] {}",taskName,addr);

            PushItemRequest pushItemRequest = new PushItemRequest();
            pushItemRequest.setItems(items);
            pushItemRequest.setTaskName(taskName);

            Channel channel = this.channelConnectManager.getChannel(addr);
            RequestMessage request  = RequestMessageBuilder.builder()
                    .action(RequestAction.PUSH_TASK)
                    .body(pushItemRequest)
                    .build();
            try {
                LeaderController.this.nettyServer.async(channel,request,timeout,future -> {
                    if(future.isSendSuccess()){
                        ResponseMessage response = future.getResponse();
                        if (response.isSuccess()) {
                            for (TaskItem item : items) {
                                context.commitStatus(item.getId(), TaskProcessStatus.SUCCESS);
                            }
                        }
                        else {
                            for (TaskItem item : items) {
                                context.commitStatus(item.getId(), TaskProcessStatus.FAILED);
                            }
                        }
                    }
                    else {
                        for (TaskItem item : items) {
                            context.commitStatus(item.getId(), TaskProcessStatus.FAILED);
                        }
                    }
                });
            } catch (Exception e) {
                log.error(e.getMessage(),e);
                BootstrapEventHolder.getInstance().publishEvent(new BootstrapEvent(addr, BootstrapEventType.CHANNEL_CLOSE));
                for (TaskItem item : items) {
                    context.commitStatus(item.getId(), TaskProcessStatus.FAILED);
                }
            }
        }
    }

}
